"""
宏观经济数据采集器模块

负责通过akshare API获取宏观经济数据，进行清洗，并保存到数据库
"""
import time
import datetime
import logging
import akshare as ak
import pandas as pd
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.macro import MacroEconomicData, EconomicCalendar
from config.settings import DATA_COLLECTION

# 配置日志
logger = logging.getLogger("macro_collector")


class MacroCollector:
    """宏观经济数据采集类"""
    
    def __init__(self):
        """初始化采集器"""
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.batch_size = DATA_COLLECTION["batch_size"]
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def get_macro_indicators(self):
        """获取宏观经济指标列表"""
        indicators = [
            {"code": "GDP", "name": "国内生产总值", "api_func": "macro_china_gdp"},
            {"code": "CPI", "name": "居民消费价格指数", "api_func": "macro_china_cpi"},
            {"code": "PPI", "name": "工业生产者出厂价格指数", "api_func": "macro_china_ppi"},
            {"code": "M2", "name": "货币供应量M2", "api_func": "macro_china_m2_yearly"},
            {"code": "PMI", "name": "采购经理人指数", "api_func": "macro_china_pmi"},
            {"code": "EXPORT", "name": "出口金额", "api_func": "macro_china_exports_yoy"},
            {"code": "IMPORT", "name": "进口金额", "api_func": "macro_china_imports_yoy"},
            {"code": "RATE", "name": "存款基准利率", "api_func": "macro_china_deposit_rate"}
        ]
        return indicators
    
    def collect_all_indicators(self):
        """采集所有宏观经济指标"""
        results = {}
        for indicator in self.get_macro_indicators():
            try:
                if indicator["code"] == "GDP":
                    results["GDP"] = self.collect_gdp_data()
                elif indicator["code"] == "CPI":
                    results["CPI"] = self.collect_cpi_data()
                # 其他指标可以类似实现
            except Exception as e:
                logger.error(f"{indicator['name']}采集失败: {str(e)}")
        return results
    
    def collect_gdp_data(self, start_date=None, end_date=None, task_id=None):
        """
        采集GDP数据
        
        Args:
            start_date: 开始日期，默认为None
            end_date: 结束日期，默认为None
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 设置日期范围
        if not start_date:
            # 从最后一条记录的下一个季度开始
            last_record = self.db.query(MacroEconomicData).filter(
                MacroEconomicData.indicator == "GDP"
            ).order_by(MacroEconomicData.date.desc()).first()
            
            if last_record:
                start_date = (last_record.date + datetime.timedelta(days=90)).strftime('%Y%m%d')
            else:
                # 默认采集近10年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365*10)).strftime('%Y%m%d')
        
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        logger.info("开始采集GDP数据，时间范围: %s 至 %s", start_date, end_date)
        
        try:
            # 重试机制
            retry_count = 0
            df = None
            
            while retry_count < self.retry_limit:
                try:
                    # 使用akshare获取GDP数据
                    df = ak.macro_china_gdp()
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning("采集GDP数据失败，重试 %d/%d: %s", 
                                  retry_count, self.retry_limit, str(e))
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df is None or df.empty:
                logger.info("在指定时间范围内没有GDP数据")
                return 0
            
            # 清洗和转换数据
            df = self._clean_macro_data(df, "GDP")
            
            # 筛选日期范围
            df['date'] = pd.to_datetime(df['date'])
            start_date_dt = pd.to_datetime(start_date)
            end_date_dt = pd.to_datetime(end_date)
            df = df[(df['date'] >= start_date_dt) & (df['date'] <= end_date_dt)]
            
            # 保存到数据库
            inserted_count = self._save_macro_data(df)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集GDP数据失败: {str(e)}"
            logger.error(error_msg)
            return 0
    
    def collect_cpi_data(self, start_date=None, end_date=None, task_id=None):
        """
        采集CPI数据
        
        Args:
            start_date: 开始日期，默认为None
            end_date: 结束日期，默认为None
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 设置日期范围
        if not start_date:
            # 从最后一条记录的下一个月开始
            last_record = self.db.query(MacroEconomicData).filter(
                MacroEconomicData.indicator == "CPI"
            ).order_by(MacroEconomicData.date.desc()).first()
            
            if last_record:
                start_date = (last_record.date + datetime.timedelta(days=31)).strftime('%Y%m%d')
            else:
                # 默认采集近10年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365*10)).strftime('%Y%m%d')
        
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        logger.info("开始采集CPI数据，时间范围: %s 至 %s", start_date, end_date)
        
        try:
            # 重试机制
            retry_count = 0
            df = None
            
            while retry_count < self.retry_limit:
                try:
                    # 使用akshare获取CPI数据
                    df = ak.macro_china_cpi()
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning("采集CPI数据失败，重试 %d/%d: %s", 
                                  retry_count, self.retry_limit, str(e))
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df is None or df.empty:
                logger.info("在指定时间范围内没有CPI数据")
                return 0
            
            # 清洗和转换数据
            df = self._clean_macro_data(df, "CPI")
            
            # 筛选日期范围
            df['date'] = pd.to_datetime(df['date'])
            start_date_dt = pd.to_datetime(start_date)
            end_date_dt = pd.to_datetime(end_date)
            df = df[(df['date'] >= start_date_dt) & (df['date'] <= end_date_dt)]
            
            # 保存到数据库
            inserted_count = self._save_macro_data(df)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集CPI数据失败: {str(e)}"
            logger.error(error_msg)
            return 0
    
    def _clean_macro_data(self, df, indicator):
        """清洗和标准化宏观经济数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.lower() for col in df.columns]
        
        # 添加指标类型
        df['indicator'] = indicator
        
        # 处理日期
        if 'date' in df.columns:
            df['date'] = pd.to_datetime(df['date'])
        
        # 处理数值类型
        numeric_cols = ['value']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 保留需要的列
        keep_cols = [col for col in ['date', 'indicator', 'value'] if col in df.columns]
        
        return df[keep_cols]
    
    def _save_macro_data(self, df):
        """保存宏观经济数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的记录
            existing_records = {
                (record.date, record.indicator) for record in 
                self.db.query(MacroEconomicData.date, MacroEconomicData.indicator).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                date = row['date'].date()
                indicator = row['indicator']
                
                # 跳过已存在的记录
                if (date, indicator) in existing_records:
                    continue
                
                macro_data = MacroEconomicData(
                    date=date,
                    indicator=indicator,
                    value=row.get('value'),
                    update_time=datetime.datetime.now()
                )
                
                batch_records.append(macro_data)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info("成功保存 %d 条宏观经济数据", inserted_count)
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error("保存宏观经济数据失败(完整性错误): %s", str(e))
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error("保存宏观经济数据失败: %s", str(e))
            raise


# 测试函数
def test_macro_collector():
    """测试宏观经济数据采集器功能"""
    collector = MacroCollector()
    
    # 1. 测试获取宏观经济指标列表
    indicators = collector.get_macro_indicators()
    print(f"支持的宏观经济指标: {len(indicators)}")
    for indicator in indicators:
        print(f"指标: {indicator['name']} ({indicator['code']})")
    
    # 2. 测试采集GDP数据
    try:
        # 采集最近5年数据
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=365*5)).strftime('%Y%m%d')
        
        count = collector.collect_gdp_data(start_date, end_date)
        print(f"成功采集 {count} 条GDP数据")
    except Exception as e:
        print(f"采集GDP数据失败: {str(e)}")
    
    # 3. 测试采集CPI数据
    try:
        # 采集最近5年数据
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=365*5)).strftime('%Y%m%d')
        
        count = collector.collect_cpi_data(start_date, end_date)
        print(f"成功采集 {count} 条CPI数据")
    except Exception as e:
        print(f"采集CPI数据失败: {str(e)}")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_macro_collector()
